package com.zl.cms.util;


import java.io.*;
import java.sql.SQLException;

import junit.framework.Assert;

import org.dbunit.DatabaseUnitException;
import org.dbunit.database.DatabaseConnection;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.database.QueryDataSet;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatDtdDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSet;
import org.dbunit.dataset.xml.FlatXmlProducer;
import org.dbunit.operation.DatabaseOperation;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.xml.sax.InputSource;

/**
 * User: zl email:c3tc3tc3t@126.com Date: 14-1-14 Time: 下午2:19
 */
public class DbunitTestCase {
  protected static IDatabaseConnection databaseConnection = null;
  private File temlateFile = null;
  private File dtdfile = null;

  /**
   * 初始化dbunit数据库连接
   */
  @BeforeClass
  public static void init() {
    try {
      databaseConnection =
          new DatabaseConnection(DbUtil.getConnection(
              "jdbc:mysql://localhost:3306/zl_cms?sessionVariables=FOREIGN_KEY_CHECKS=0", "zl",
              "zl123"));
    } catch (DatabaseUnitException e) {
      e.printStackTrace();
    }
  }

  /**
   * 得到测试数据表的DataSet
   * 
   * @param tname
   * @return
   */
  protected IDataSet createTestDataSet(String tname) throws DatabaseUnitException, IOException {
    InputStream is = this.getClass().getClassLoader().getResourceAsStream(tname + ".xml");
    Assert.assertNotNull("the Test datafile " + tname + " not exist", is);
    return new FlatXmlDataSet(new FlatXmlProducer(new InputSource(is), new FlatDtdDataSet(
        new FileReader(dtdfile))));
  }

  /**
   * 测试之前将数据库的所有表备份
   */
  protected void backAllTables() throws SQLException, IOException, DataSetException {
    IDataSet dataSet = databaseConnection.createDataSet();
    writeTableToFile(dataSet);
  }

  /**
   * 要备份的表的数据写入back.xml
   * 
   * @param dataSet
   * @throws Exception
   */
  protected void writeTableToFile(IDataSet dataSet) throws IOException, DataSetException {
    temlateFile = File.createTempFile("back", "xml");
    dtdfile = File.createTempFile("back", "dtd");
    FlatDtdDataSet.write(dataSet, new FileWriter(dtdfile));
    FlatXmlDataSet.write(dataSet, new FileWriter(temlateFile));
  }

  /**
   * 备份指定的表
   * 
   * @param args 表名数组
   * @throws Exception
   */
  protected void backCustomeTables(String[] args) throws DataSetException, IOException {
    QueryDataSet queryDataSet = new QueryDataSet(databaseConnection);
    for (String arg : args) {
      queryDataSet.addTable(arg);
    }

    writeTableToFile(queryDataSet);
  }

  /**
   * 备份一张指定的表数据
   * 
   * @param tname
   * @throws Exception
   */
  protected void backOneTable(String tname) throws IOException, DataSetException {
    backCustomeTables(new String[] {tname});
  }


  /**
   * 还原测试之前的数据到数据库中
   * 
   * @throws Exception
   */
  protected void resumeTable() throws DatabaseUnitException, SQLException, IOException {
    IDataSet ds =
        new FlatXmlDataSet(new FlatXmlProducer(new InputSource(new FileInputStream(temlateFile)),
            new FlatDtdDataSet(new FileReader(dtdfile))));
    DatabaseOperation.CLEAN_INSERT.execute(databaseConnection, ds);

  }

  /**
   * 关闭测试dbunit数据库链接
   */
  @AfterClass
  public static void destory() {
    try {
      if (null != databaseConnection) {
        databaseConnection.close();
      }
    } catch (SQLException e) {
      e.printStackTrace();
    }

  }
}
